#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
XML到Excel数据处理工具 v1.0.0

此工具用于解析XML脆弱性函数数据并将其转换为Excel格式，
支持数据验证、结构化处理和批量导出。

Copyright (c) 2023 XML数据分析工具开发团队
License: MIT
"""

import os
import sys
import json
import glob
import xml.etree.ElementTree as ET
import pandas as pd
import numpy as np
from typing import Dict, List, Tuple, Optional, Union
import logging
from pathlib import Path


class XMLDataProcessor:
    """
    XML数据处理器，专注于脆弱性函数数据的解析与转换
    
    此类提供XML数据解析、DataFrame构建和Excel导出的完整工作流
    """
    
    def __init__(self, log_level=logging.INFO):
        """
        初始化XML数据处理器
        
        Args:
            log_level: 日志级别，默认为INFO
        """
        # 配置日志
        self._setup_logging(log_level)
        self.logger.info("XML数据处理器初始化完成")
    
    def _setup_logging(self, log_level):
        """
        设置日志系统
        
        Args:
            log_level: 日志级别
        """
        self.logger = logging.getLogger("XMLDataProcessor")
        self.logger.setLevel(log_level)
        
        # 如果没有处理器则添加一个
        if not self.logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            handler.setFormatter(formatter)
            self.logger.addHandler(handler)
    
    def parse_xml(self, xml_path: str) -> List[Dict]:
        """
        解析XML文件，提取脆弱性函数数据
        
        Args:
            xml_path: XML文件路径
            
        Returns:
            包含所有脆弱性函数数据的列表
        
        Raises:
            FileNotFoundError: 文件不存在
            ET.ParseError: XML解析错误
        """
        self.logger.info(f"开始解析XML文件: {xml_path}")
        
        if not os.path.exists(xml_path):
            self.logger.error(f"文件不存在: {xml_path}")
            raise FileNotFoundError(f"文件不存在: {xml_path}")
        
        try:
            # 使用ElementTree解析XML
            tree = ET.parse(xml_path)
            root = tree.getroot()
            
            # 提取XML命名空间
            ns = ''
            if '}' in root.tag:
                ns = root.tag.split('}')[0].strip('{')
                
            # 准备数据存储结构
            vulnerability_data = []
            
            # 定位脆弱性模型
            vulnerability_model = None
            if ns:
                nsmap = {'nrml': ns}
                vulnerability_model = root.find('.//{%s}vulnerabilityModel' % ns)
            else:
                vulnerability_model = root.find('.//vulnerabilityModel')
            
            if vulnerability_model is None:
                self.logger.warning("未找到vulnerabilityModel节点")
                return []
            
            # 提取模型属性
            model_id = vulnerability_model.get('id', '')
            asset_category = vulnerability_model.get('assetCategory', '')
            loss_category = vulnerability_model.get('lossCategory', '')
            
            # 遍历所有脆弱性函数
            vuln_funcs = []
            if ns:
                vuln_funcs = vulnerability_model.findall('.//{%s}vulnerabilityFunction' % ns)
            else:
                vuln_funcs = vulnerability_model.findall('.//vulnerabilityFunction')
                
            for vuln_func in vuln_funcs:
                try:
                    # 提取脆弱性函数属性
                    vuln_id = vuln_func.get('id', '')
                    dist = vuln_func.get('dist', '')
                    
                    # 解析imls元素及其属性
                    imls_elem = None
                    if ns:
                        imls_elem = vuln_func.find('.//{%s}imls' % ns)
                    else:
                        imls_elem = vuln_func.find('.//imls')
                        
                    if imls_elem is None:
                        self.logger.warning(f"函数 {vuln_id} 中未找到imls元素，跳过")
                        continue
                    
                    imt = imls_elem.get('imt', '')
                    
                    # 提取数值序列
                    imls_text = imls_elem.text.strip() if imls_elem.text else ""
                    imls = [float(x) for x in imls_text.split()] if imls_text else []
                    
                    # 解析meanLRs元素
                    meanLRs_elem = None
                    if ns:
                        meanLRs_elem = vuln_func.find('.//{%s}meanLRs' % ns)
                    else:
                        meanLRs_elem = vuln_func.find('.//meanLRs')
                        
                    if meanLRs_elem is None:
                        self.logger.warning(f"函数 {vuln_id} 中未找到meanLRs元素，跳过")
                        continue
                    
                    meanLRs_text = meanLRs_elem.text.strip() if meanLRs_elem.text else ""
                    meanLRs = [float(x) for x in meanLRs_text.split()] if meanLRs_text else []
                    
                    # 解析covLRs元素
                    covLRs_elem = None
                    if ns:
                        covLRs_elem = vuln_func.find('.//{%s}covLRs' % ns)
                    else:
                        covLRs_elem = vuln_func.find('.//covLRs')
                        
                    if covLRs_elem is None:
                        self.logger.warning(f"函数 {vuln_id} 中未找到covLRs元素，跳过")
                        continue
                    
                    covLRs_text = covLRs_elem.text.strip() if covLRs_elem.text else ""
                    covLRs = [float(x) for x in covLRs_text.split()] if covLRs_text else []
                    
                    # 验证数据一致性
                    if not self._validate_data_consistency(imls, meanLRs, covLRs):
                        self.logger.warning(f"函数 {vuln_id} 数据长度不一致，跳过")
                        continue
                    
                    # 构建数据字典
                    vuln_data = {
                        'model_id': model_id,
                        'asset_category': asset_category,
                        'loss_category': loss_category,
                        'vulnerability_id': vuln_id,
                        'dist': dist,
                        'imt': imt,
                        'imls': imls,
                        'meanLRs': meanLRs,
                        'covLRs': covLRs
                    }
                    
                    vulnerability_data.append(vuln_data)
                    self.logger.debug(f"成功解析脆弱性函数: {vuln_id}")
                    
                except Exception as e:
                    self.logger.error(f"解析脆弱性函数时出错: {e}")
                    continue
            
            self.logger.info(f"成功解析 {len(vulnerability_data)} 个脆弱性函数")
            return vulnerability_data
            
        except ET.ParseError as e:
            self.logger.error(f"XML解析错误: {e}")
            raise
        except Exception as e:
            self.logger.error(f"解析过程中发生错误: {e}")
            raise
    
    def _validate_data_consistency(self, imls: List[float], meanLRs: List[float], 
                                  covLRs: List[float]) -> bool:
        """
        验证数据一致性，确保三列数据长度相同
        
        Args:
            imls: 强度测量水平数据
            meanLRs: 平均损失率数据
            covLRs: 变异系数数据
            
        Returns:
            数据是否一致
        """
        lengths = [len(imls), len(meanLRs), len(covLRs)]
        if len(set(lengths)) != 1:
            self.logger.warning(f"数据长度不一致: imls={lengths[0]}, meanLRs={lengths[1]}, covLRs={lengths[2]}")
            return False
        
        if any(l == 0 for l in lengths):
            self.logger.warning("至少有一个数据序列为空")
            return False
            
        return True
    
    def create_dataframe(self, vulnerability_data: List[Dict]) -> pd.DataFrame:
        """
        将脆弱性函数数据转换为DataFrame
        
        Args:
            vulnerability_data: 脆弱性函数数据列表
            
        Returns:
            包含所有数据的DataFrame
        """
        self.logger.info("开始构建DataFrame")
        
        if not vulnerability_data:
            self.logger.warning("没有可用的脆弱性函数数据")
            return pd.DataFrame()
        
        # 准备存储展开数据的列表
        rows = []
        
        for vuln_data in vulnerability_data:
            # 提取基本属性
            base_data = {
                'model_id': vuln_data['model_id'],
                'asset_category': vuln_data['asset_category'],
                'loss_category': vuln_data['loss_category'],
                'vulnerability_id': vuln_data['vulnerability_id'],
                'dist': vuln_data['dist'],
                'imt': vuln_data['imt']
            }
            
            # 获取数据序列
            imls = vuln_data['imls']
            meanLRs = vuln_data['meanLRs']
            covLRs = vuln_data['covLRs']
            
            # 将每组对应的值合并为一行
            for i in range(len(imls)):
                row_data = base_data.copy()
                row_data['imls'] = imls[i]
                row_data['meanLRs'] = meanLRs[i]
                row_data['covLRs'] = covLRs[i]
                rows.append(row_data)
        
        # 创建DataFrame
        df = pd.DataFrame(rows)
        
        # 转换数值列的类型
        numeric_cols = ['imls', 'meanLRs', 'covLRs']
        for col in numeric_cols:
            df[col] = pd.to_numeric(df[col], errors='coerce')
        
        self.logger.info(f"成功创建DataFrame，包含 {len(df)} 行数据")
        return df
    
    def export_to_excel(self, df: pd.DataFrame, output_path: str) -> bool:
        """
        将DataFrame导出为Excel文件
        
        Args:
            df: 包含脆弱性函数数据的DataFrame
            output_path: 输出Excel文件路径
            
        Returns:
            是否成功导出
            
        Raises:
            PermissionError: 文件权限错误
        """
        self.logger.info(f"开始导出数据到Excel: {output_path}")
        
        if df.empty:
            self.logger.warning("DataFrame为空，无数据可导出")
            return False
        
        try:
            # 创建输出目录（如果不存在）
            output_dir = os.path.dirname(output_path)
            if output_dir and not os.path.exists(output_dir):
                os.makedirs(output_dir)
                
            # 使用openpyxl引擎导出到Excel
            with pd.ExcelWriter(output_path, engine='openpyxl') as writer:
                # 按vulnerability_id分组创建工作表
                for vuln_id, group_df in df.groupby('vulnerability_id'):
                    # 清理工作表名称（Excel工作表名称有限制）
                    sheet_name = self._clean_sheet_name(vuln_id)
                    group_df.to_excel(writer, sheet_name=sheet_name, index=False)
                
                # 创建一个汇总表
                summary_df = df.drop_duplicates(subset=['vulnerability_id'])
                summary_cols = ['vulnerability_id', 'dist', 'imt', 'model_id', 
                               'asset_category', 'loss_category']
                summary_df = summary_df[summary_cols]
                summary_df.to_excel(writer, sheet_name='汇总', index=False)
            
            self.logger.info(f"成功导出数据到Excel: {output_path}")
            return True
            
        except PermissionError:
            self.logger.error(f"无权限写入文件: {output_path}")
            raise
        except Exception as e:
            self.logger.error(f"导出Excel时出错: {e}")
            raise
    
    def _clean_sheet_name(self, name: str) -> str:
        """
        清理工作表名称，确保符合Excel要求
        
        Args:
            name: 原始工作表名称
            
        Returns:
            清理后的工作表名称
        """
        # Excel工作表名称不能超过31个字符
        if len(name) > 31:
            name = name[:31]
        
        # 移除Excel不允许的字符
        forbidden_chars = [':', '\\', '/', '?', '*', '[', ']']
        for char in forbidden_chars:
            name = name.replace(char, '_')
            
        return name
    
    def create_summary_excel(self, excel_path: str, output_path: str = None) -> bool:
        """
        从现有Excel文件创建汇总表格
        
        读取Excel中的所有工作表，提取imls和meanLRs数据，
        创建一个新的汇总表，以imls为第一列，各工作表的meanLRs为后续列
        
        Args:
            excel_path: 输入Excel文件路径
            output_path: 输出Excel文件路径（默认为输入文件名加_renew后缀）
            
        Returns:
            是否成功创建汇总表
        """
        self.logger.info(f"开始创建汇总表格: {excel_path}")
        
        try:
            # 确定输出文件路径
            if not output_path:
                base_name = os.path.splitext(excel_path)[0]
                output_path = f"{base_name}_renew.xlsx"
                
            # 检查文件是否存在
            if not os.path.exists(excel_path):
                self.logger.error(f"文件不存在: {excel_path}")
                return False
                
            # 读取Excel文件中的所有工作表
            excel_data = pd.read_excel(excel_path, sheet_name=None)
            
            # 排除"汇总"工作表
            if "汇总" in excel_data:
                del excel_data["汇总"]
                
            if not excel_data:
                self.logger.warning("Excel文件中没有有效的工作表")
                return False
                
            # 创建汇总DataFrame
            summary_data = {}
            sheet_names = list(excel_data.keys())
            
            # 获取第一个工作表的imls数据作为基准
            first_sheet = excel_data[sheet_names[0]]
            if "imls" not in first_sheet.columns:
                self.logger.error("工作表中未找到imls列")
                return False
                
            summary_data["imls"] = first_sheet["imls"].values
            
            # 提取每个工作表的meanLRs数据
            for sheet_name, sheet_df in excel_data.items():
                if "meanLRs" not in sheet_df.columns:
                    self.logger.warning(f"工作表 {sheet_name} 中未找到meanLRs列，跳过")
                    continue
                    
                summary_data[sheet_name] = sheet_df["meanLRs"].values
                
            # 创建汇总DataFrame
            summary_df = pd.DataFrame(summary_data)
            
            # 导出到Excel
            with pd.ExcelWriter(output_path, engine='openpyxl') as writer:
                summary_df.to_excel(writer, sheet_name='汇总数据', index=False)
                
            self.logger.info(f"成功创建汇总表格: {output_path}")
            return True
            
        except Exception as e:
            self.logger.error(f"创建汇总表格时出错: {e}")
            return False
    
    def process_xml_to_excel(self, xml_path: str, output_path: str) -> bool:
        """
        完整处理流程：解析XML并导出到Excel
        
        Args:
            xml_path: XML文件路径
            output_path: 输出Excel文件路径
            
        Returns:
            处理是否成功
        """
        try:
            self.logger.info(f"开始处理: {xml_path} -> {output_path}")
            
            # 解析XML
            vulnerability_data = self.parse_xml(xml_path)
            
            if not vulnerability_data:
                self.logger.warning("未找到脆弱性函数数据")
                return False
            
            # 创建DataFrame
            df = self.create_dataframe(vulnerability_data)
            
            if df.empty:
                self.logger.warning("生成的DataFrame为空")
                return False
            
            # 导出到Excel
            result = self.export_to_excel(df, output_path)
            
            self.logger.info(f"处理{'成功' if result else '失败'}")
            return result
            
        except Exception as e:
            self.logger.error(f"处理过程中发生错误: {e}")
            return False

    def process_directory(self, input_dir: str, output_dir: str, pattern: str = "*.xml", create_summary: bool = True) -> Dict[str, bool]:
        """
        处理目录中的所有XML文件
        
        Args:
            input_dir: 输入目录路径
            output_dir: 输出目录路径
            pattern: 文件匹配模式
            create_summary: 是否为每个文件创建汇总表格
            
        Returns:
            处理结果字典，键为文件名，值为处理是否成功
        """
        self.logger.info(f"开始批量处理目录: {input_dir} -> {output_dir}")
        
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        
        # 查找所有匹配的XML文件
        xml_files = glob.glob(os.path.join(input_dir, pattern))
        
        if not xml_files:
            self.logger.warning(f"在目录 {input_dir} 中没有找到匹配 {pattern} 的文件")
            return {}
            
        self.logger.info(f"找到 {len(xml_files)} 个XML文件待处理")
        
        # 处理每个文件
        results = {}
        for xml_file in xml_files:
            file_name = os.path.basename(xml_file)
            base_name = os.path.splitext(file_name)[0]
            excel_output = os.path.join(output_dir, f"{base_name}.xlsx")
            
            self.logger.info(f"处理文件: {file_name}")
            
            try:
                # 处理XML文件
                success = self.process_xml_to_excel(xml_file, excel_output)
                
                # 创建汇总表格
                if success and create_summary:
                    summary_success = self.create_summary_excel(excel_output)
                    if not summary_success:
                        self.logger.warning(f"为文件 {file_name} 创建汇总表格失败")
                
                results[file_name] = success
                self.logger.info(f"文件 {file_name} 处理{'成功' if success else '失败'}")
                
            except Exception as e:
                self.logger.error(f"处理文件 {file_name} 时出错: {e}")
                results[file_name] = False
                
        # 统计处理结果
        success_count = sum(1 for result in results.values() if result)
        self.logger.info(f"批量处理完成: 共 {len(results)} 个文件，{success_count} 个成功，{len(results) - success_count} 个失败")
        
        return results


def load_config(config_path="config.json") -> Dict:
    """
    加载配置文件
    
    Args:
        config_path: 配置文件路径
        
    Returns:
        配置字典
    """
    if os.path.exists(config_path):
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"加载配置文件出错: {e}")
    return {}


def main():
    """
    主函数，处理命令行参数并执行转换
    """
    # 加载配置文件
    config = load_config()
    
    # 配置命令行参数
    import argparse
    parser = argparse.ArgumentParser(description="XML数据处理工具 - 将XML脆弱性函数数据转换为Excel格式")
    parser.add_argument("xml_file", nargs="?", 
                       help="输入XML文件路径（单文件模式）")
    parser.add_argument("-o", "--output", help="输出Excel文件路径（单文件模式）")
    parser.add_argument("-i", "--input-dir", help="输入目录路径（批处理模式）")
    parser.add_argument("-d", "--output-dir", help="输出目录路径（批处理模式）")
    parser.add_argument("-p", "--pattern", default="*.xml", help="文件匹配模式（批处理模式）")
    parser.add_argument("-v", "--verbose", action="store_true", help="显示详细日志")
    parser.add_argument("-c", "--config", help="指定配置文件路径")
    parser.add_argument("--no-summary", action="store_true", help="不创建汇总表格")
    parser.add_argument("--mode", choices=["single", "batch"], help="处理模式：单文件或批处理")
    args = parser.parse_args()
    
    # 如果指定了配置文件，重新加载
    if args.config:
        config = load_config(args.config)
    
    # 设置日志级别
    log_level = logging.DEBUG if args.verbose or config.get("verbose", False) else logging.INFO
    
    # 创建处理器实例
    processor = XMLDataProcessor(log_level=log_level)
    
    # 确定处理模式
    mode = args.mode or config.get("mode", "single")
    
    if mode == "batch":
        # 批处理模式
        input_dir = args.input_dir or config.get("batch_mode", {}).get("input_directory", "./")
        output_dir = args.output_dir or config.get("batch_mode", {}).get("output_directory", "./output")
        pattern = args.pattern or config.get("batch_mode", {}).get("file_pattern", "*.xml")
        create_summary = not args.no_summary and config.get("batch_mode", {}).get("create_summary", True)
        
        # 执行批处理
        results = processor.process_directory(input_dir, output_dir, pattern, create_summary)
        
        # 设置退出代码
        success = any(results.values()) if results else False
        sys.exit(0 if success else 1)
        
    else:
        # 单文件模式
        single_config = config.get("single_mode", {})
        xml_file = args.xml_file or single_config.get("xml_file_path")
        
        if not xml_file:
            print("错误: 未指定XML文件路径")
            parser.print_help()
            sys.exit(1)
            
        output = args.output or single_config.get("output_excel_path")
        
        if not output:
            base_name = os.path.splitext(xml_file)[0]
            output = f"{base_name}.xlsx"
        
        # 处理单个文件
        success = processor.process_xml_to_excel(xml_file, output)
        
        # 创建汇总表格
        if success and not args.no_summary:
            processor.create_summary_excel(output)
        
        # 设置退出代码
        sys.exit(0 if success else 1)


if __name__ == "__main__":
    main() 